from argparse import Namespace
import random
from typing import Callable, List, Union, Tuple

import numpy as np
import torch
from torch.utils.data.dataset import Dataset
from rdkit import Chem

from .scaler import StandardScaler
from .mol_tree import MolTree, MolTreeNode
from features import get_features_generator
import copy
import random
import pickle
from mol2vec.features import mol2alt_sentence

def load_pickle(filepath):
    obj = pickle.load(open(filepath, 'rb'))
    return obj


smiles2moltree = None


class MoleculeDatapoint:
    """A MoleculeDatapoint contains a single molecule and its associated features and targets."""

    smiles2moltree = None

    def __init__(self,
                 line: List[str],
                 args: Namespace = None,
                 features: np.ndarray = None,
                 use_compound_names: bool = False,
                 label_index=-1):
        """
        Initializes a MoleculeDatapoint, which contains a single molecule.

        :param line: A list of strings generated by separating a line in a data CSV file by comma.
        :param args: Arguments.
        :param features: A numpy array containing additional features (ex. Morgan fingerprint).
        :param use_compound_names: Whether the data CSV includes the compound name on each line.
        """
        if args is not None:
            self.features_generator = args.features_generator
            self.args = args
        else:
            self.features_generator = self.args = None

        if features is not None and self.features_generator is not None:
            raise ValueError('Currently cannot provide both loaded features and a features generator.')

        self.features = features

        if use_compound_names:
            self.compound_name = line[0]  # str
            line = line[1:]
        else:
            self.compound_name = None

        self.smiles = line[0]  # str
        self.mol = Chem.MolFromSmiles(self.smiles)

        # Generate additional features if given a generator
        if self.features_generator is not None:
            self.features = []

            for fg in self.features_generator:
                features_generator = get_features_generator(fg)
                if self.mol is not None and self.mol.GetNumHeavyAtoms() > 0:
                    self.features.extend(features_generator(self.mol))

            self.features = np.array(self.features)

        # Fix nans in features
        if self.features is not None:
            replace_token = 0
            self.features = np.where(np.isnan(self.features), replace_token, self.features)

        # Create targets
        self.targets = [float(x) if x != '' else None for x in line[1:]]

        # for junction tree prediction
        self.mol_tree = self.get_mol_tree() if self.args.jt else None

    def set_features(self, features: np.ndarray):
        """
        Sets the features of the molecule.

        :param features: A 1-D numpy array of features for the molecule.
        """
        self.features = features

    def num_tasks(self) -> int:
        """
        Returns the number of prediction tasks.

        :return: The number of tasks.
        """
        return len(self.targets)

    def set_targets(self, targets: List[float]):
        """
        Sets the targets of a molecule.

        :param targets: A list of floats containing the targets.
        """
        self.targets = targets

    def get_mol_tree(self) -> MolTree:
        assert self.args.jt == True
        global smiles2moltree
        if smiles2moltree is None:
            smiles2moltree = load_pickle(self.args.smiles2moltree_file)
        return smiles2moltree[self.smiles]


class MoleculeDataset(Dataset):
    """A MoleculeDataset contains a list of molecules and their associated features and targets."""

    def __init__(self, data: List[MoleculeDatapoint]):
        """
        Initializes a MoleculeDataset, which contains a list of MoleculeDatapoints (i.e. a list of molecules).

        :param data: A list of MoleculeDatapoints.
        """
        self.data = data
        self.args = self.data[0].args if len(self.data) > 0 else None
        self.scaler = None

    def compound_names(self) -> List[str]:
        """
        Returns the compound names associated with the molecule (if they exist).

        :return: A list of compound names or None if the dataset does not contain compound names.
        """
        if len(self.data) == 0 or self.data[0].compound_name is None:
            return None

        return [d.compound_name for d in self.data]

    def smiles(self) -> List[str]:
        """
        Returns the smiles strings associated with the molecules.

        :return: A list of smiles strings.
        """
        return [d.smiles for d in self.data]
    
    def mols(self) -> List[Chem.Mol]:
        """
        Returns the RDKit molecules associated with the molecules.

        :return: A list of RDKit Mols.
        """
        return [d.mol for d in self.data]

    # add for jt graph model
    def mol_trees(self) -> List[MolTree]:
        assert self.args.jt == True
        return [d.mol_tree for d in self.data if d.mol_tree.has_root()]

    def features(self) -> List[np.ndarray]:
        """
        Returns the features associated with each molecule (if they exist).

        :return: A list of 1D numpy arrays containing the features for each molecule or None if there are no features.
        """
        if len(self.data) == 0 or self.data[0].features is None:
            return None

        return [d.features for d in self.data]

    def targets(self) -> List[List[float]]:
        """
        Returns the targets associated with each molecule.

        :return: A list of lists of floats containing the targets.
        """
        # add for jt graph model
        if not self.args.jt:
            return [d.targets for d in self.data]
        else:
            return [d.targets for d in self.data if d.mol_tree.has_root()]

    def num_tasks(self) -> int:
        """
        Returns the number of prediction tasks.

        :return: The number of tasks.
        """
        return self.data[0].num_tasks() if len(self.data) > 0 else None

    def features_size(self) -> int:
        """
        Returns the size of the features array associated with each molecule.

        :return: The size of the features.
        """
        return len(self.data[0].features) if len(self.data) > 0 and self.data[0].features is not None else None

    def shuffle(self, seed: int = None):
        """
        Shuffles the dataset.

        :param seed: Optional random seed.
        """
        if seed is not None:
            random.seed(seed)
        random.shuffle(self.data)
    
    def normalize_features(self, scaler: StandardScaler = None, replace_nan_token: int = 0) -> StandardScaler:
        """
        Normalizes the features of the dataset using a StandardScaler (subtract mean, divide by standard deviation).

        If a scaler is provided, uses that scaler to perform the normalization. Otherwise fits a scaler to the
        features in the dataset and then performs the normalization.

        :param scaler: A fitted StandardScaler. Used if provided. Otherwise a StandardScaler is fit on
        this dataset and is then used.
        :param replace_nan_token: What to replace nans with.
        :return: A fitted StandardScaler. If a scaler is provided, this is the same scaler. Otherwise, this is
        a scaler fit on this dataset.
        """
        if len(self.data) == 0 or self.data[0].features is None:
            return None

        if scaler is not None:
            self.scaler = scaler

        elif self.scaler is None:
            features = np.vstack([d.features for d in self.data])
            self.scaler = StandardScaler(replace_nan_token=replace_nan_token)
            self.scaler.fit(features)

        for d in self.data:
            d.set_features(self.scaler.transform(d.features.reshape(1, -1))[0])

        return self.scaler
    
    def set_targets(self, targets: List[List[float]]):
        """
        Sets the targets for each molecule in the dataset. Assumes the targets are aligned with the datapoints.

        :param targets: A list of lists of floats containing targets for each molecule. This must be the
        same length as the underlying dataset.
        """
        assert len(self.data) == len(targets)
        for i in range(len(self.data)):
            self.data[i].set_targets(targets[i])

    def sort(self, key: Callable):
        """
        Sorts the dataset using the provided key.

        :param key: A function on a MoleculeDatapoint to determine the sorting order.
        """
        self.data.sort(key=key)

    def __len__(self) -> int:
        """
        Returns the length of the dataset (i.e. the number of molecules).

        :return: The length of the dataset.
        """
        return len(self.data)

    def __getitem__(self, item) -> Union[MoleculeDatapoint, List[MoleculeDatapoint], dict]:
        """
        Gets one or more MoleculeDatapoints via an index or slice.

        :param item: An index (int) or a slice object.
        :return: A MoleculeDatapoint if an int is provided or a list of MoleculeDatapoints if a slice is provided.
        """
        return self.data[item]

class DDIDatapoint:
    """A MoleculeDatapoint contains a single ddi and its associated features and targets."""

    def __init__(self,
                 dictionary: dict,
                 args: Namespace = None,
                 features_1: np.ndarray = None,
                 features_2: np.ndarray = None,
                 ):
        """
        Initializes a MoleculeDatapoint, which contains a single molecule.

        :param line: A list of strings generated by separating a line in a data CSV file by comma.
        :param args: Arguments.
        :param features: A numpy array containing additional features (ex. Morgan fingerprint).
        :param use_compound_names: Whether the data CSV includes the compound name on each line.
        """
        if args is not None:
            self.features_generator = args.features_generator
            self.args = args
        else:
            self.features_generator = self.args = None

        if features_1 is not None and self.features_generator is not None:
            raise ValueError('Currently cannot provide both loaded features and a features generator.')
        if features_2 is not None and self.features_generator is not None:
            raise ValueError('Currently cannot provide both loaded features and a features generator.')

        self.features_1 = features_1
        self.features_2 = features_2

        self.smiles_1 = dictionary['smiles_1']
        self.smiles_2 = dictionary['smiles_2']

        self.mol_1 = Chem.MolFromSmiles(self.smiles_1)
        self.mol_2 = Chem.MolFromSmiles(self.smiles_2)

        # Generate additional features if given a generator
        if self.features_generator is not None:
            self.features_1 = []
            self.features_2 = []

            for fg in self.features_generator:
                features_generator = get_features_generator(fg)
                if self.mol_1 is not None and self.mol_1.GetNumHeavyAtoms() > 0:
                    self.features_1.extend(features_generator(self.smiles_1))
                if self.mol_2 is not None and self.mol_2.GetNumHeavyAtoms() > 0:
                    self.features_2.extend(features_generator(self.smiles_2))

            self.features_1 = np.array(self.features_1)
            self.features_2 = np.array(self.features_2)

        # Fix nans in features
        if self.features_1 is not None:
            replace_token = 0
            self.features_1 = np.where(np.isnan(self.features_1), replace_token, self.features_1)
        # Fix nans in features
        if self.features_2 is not None:
            replace_token = 0
            self.features_2 = np.where(np.isnan(self.features_2), replace_token, self.features_2)

        # Create targets
        if args.num_labels is not None:
            target_scalar = int(dictionary['label'])
            self.targets = [target_scalar]
        else:
            target_scalar = int(dictionary['label'])
            self.targets = [target_scalar]

        # for junction tree prediction
        self.mol_tree_pair = self.get_mol_tree_pair() if self.args.jt else None

    def set_features(self, features_1: np.ndarray, features_2: np.ndarray):
        """
        Sets the features of the molecule.

        :param features: A 1-D numpy array of features for the molecule.
        """
        self.features_1 = features_1
        self.features_2 = features_2

    def num_tasks(self) -> int:
        """
        Returns the number of prediction tasks.

        :return: The number of tasks.
        """
        # if self.args.dataset_type == 'multilabel':
        #     return self.args.num_labels
        return len(self.targets)

    def set_targets(self, targets: List[float]):
        """
        Sets the targets of a molecule.

        :param targets: A list of floats containing the targets.
        """
        self.targets = targets

    def mirror_features(self):
        temp = self.features_1
        self.features_1 = self.features_2
        self.features_2 = temp

    def mirror_smiles_pair(self):
        temp = self.smiles_1
        self.smiles_1 = self.smiles_2
        self.smiles_2 = temp

    def mirror_mol_pair(self):
        temp = self.mol_1
        self.mol_1 = self.mol_2
        self.mol_2 = temp

    # for jt graph model
    def mirror_mol_tree_pair(self):
        temp = self.mol_tree_pair[0]
        self.mol_tree_pair[0] = self.mol_tree_pair[1]
        self.mol_tree_pair[1] = temp

    def mirror(self):
        ano_datapoint = copy.deepcopy(self)
        ano_datapoint.mirror_features()
        ano_datapoint.mirror_smiles_pair()
        ano_datapoint.mirror_mol_pair()
        # for jt graph mdoel
        if self.args.jt:
            ano_datapoint.mirror_mol_tree_pair()
        return ano_datapoint

    # for jt graph model
    def get_mol_tree_pair(self) -> Tuple[MolTree, MolTree]:
        assert self.args.jt == True
        global smiles2moltree
        if smiles2moltree is None:
            smiles2moltree = load_pickle(self.args.smiles2moltree_file)
        return (smiles2moltree[self.smiles_1], smiles2moltree[self.smiles_2])


class DDIDataset(Dataset):
    """A MoleculeDataset contains a list of molecules and their associated features and targets."""

    def __init__(self, data: List[DDIDatapoint]):
        """
        Initializes a MoleculeDataset, which contains a list of MoleculeDatapoints (i.e. a list of molecules).

        :param data: A list of MoleculeDatapoints.
        """
        self.data = data
        self.args = self.data[0].args if len(self.data) > 0 else None
        self.scaler = None

    def augment(self):
        data_mirror = [d.mirror() for d in self.data]
        data_aug = copy.deepcopy(self.data)
        data_aug.extend(data_mirror)
        random.shuffle(data_aug)
        assert len(data_aug) == 2 * len(self.data)
        return DDIDataset(data_aug)

    def mirror(self):
        data_mirror = [d.mirror() for d in self.data]
        return DDIDataset(data_mirror)

    def smiles_pairs(self):
        """
        Returns the smiles strings associated with the molecules.

        :return: A list of smiles strings.
        """
        return [(d.smiles_1, d.smiles_2) for d in self.data]

    def mol_pairs(self):
        """
        Returns the RDKit molecules associated with the molecules.

        :return: A list of RDKit Mols.
        """
        return [(d.mol_1, d.mol_2) for d in self.data]

    # add for jt graph model
    def mol_tree_pairs(self) -> List[Tuple[MolTree, MolTree]]:
        assert self.args.jt == True
        return [d.mol_tree_pair for d in self.data
                if d.mol_tree_pair[0].has_root() \
                and d.mol_tree_pair[1].has_root()
                ]

    def features_pairs(self):
        """
        Returns the features associated with each molecule (if they exist).

        :return: A list of 1D numpy arrays containing the features for each molecule or None if there are no features.
        """
        if len(self.data) == 0 or self.data[0].features_1 is None or self.data[0].features_2 is None:
            return None

        return [(d.features_1, d.features_2) for d in self.data]

    def targets(self) -> List[List[int]]:
        """
        Returns the targets associated with each molecule.

        :return: A list of lists of floats containing the targets.
        """
        if not self.args.jt:
            return [d.targets for d in self.data]
        # for jt graph model
        else:
            return [d.targets for d in self.data if d.mol_tree_pair[0].has_root() and d.mol_tree_pair[1].has_root()]

    def num_tasks(self) -> int:
        """
        Returns the number of prediction tasks.

        :return: The number of tasks.
        """
        return self.data[0].num_tasks() if len(self.data) > 0 else None

    def features_size(self) -> int:
        """
        Returns the size of the features array associated with each molecule.

        :return: The size of the features.
        """
        return len(self.data[0].features_1) if len(self.data) > 0 and self.data[0].features_1 is not None else None

    def shuffle(self, seed: int = None):
        """
        Shuffles the dataset.

        :param seed: Optional random seed.
        """
        if seed is not None:
            random.seed(seed)
        random.shuffle(self.data)

    def normalize_features(self, scaler: StandardScaler = None, replace_nan_token: int = 0) -> StandardScaler:
        """
        Normalizes the features of the dataset using a StandardScaler (subtract mean, divide by standard deviation).

        If a scaler is provided, uses that scaler to perform the normalization. Otherwise fits a scaler to the
        features in the dataset and then performs the normalization.

        :param scaler: A fitted StandardScaler. Used if provided. Otherwise a StandardScaler is fit on
        this dataset and is then used.
        :param replace_nan_token: What to replace nans with.
        :return: A fitted StandardScaler. If a scaler is provided, this is the same scaler. Otherwise, this is
        a scaler fit on this dataset.
        """
        if len(self.data) == 0 or self.data[0].features_1 is None or self.data[0].features_2 is None:
            return None

        if scaler is not None:
            self.scaler = scaler

        elif self.scaler is None:
            features_1 = np.vstack([d.features_1 for d in self.data])
            self.scaler = StandardScaler(replace_nan_token=replace_nan_token)
            self.scaler.fit(features_1)

            features_2 = np.vstack([d.features_2 for d in self.data])
            self.scaler = StandardScaler(replace_nan_token=replace_nan_token)
            self.scaler.fit(features_2)

        for d in self.data:
            d.set_features(self.scaler.transform(d.features_1.reshape(1, -1))[0],
                           self.scaler.transform(d.features_2.reshape(1, -1)[0]))

        return self.scaler

    def set_targets(self, targets: List[List[int]]):
        """
        Sets the targets for each molecule in the dataset. Assumes the targets are aligned with the datapoints.

        :param targets: A list of lists of floats containing targets for each molecule. This must be the
        same length as the underlying dataset.
        """
        assert len(self.data) == len(targets)
        for i in range(len(self.data)):
            self.data[i].set_targets(targets[i])

    def sort(self, key: Callable):
        """
        Sorts the dataset using the provided key.

        :param key: A function on a MoleculeDatapoint to determine the sorting order.
        """
        self.data.sort(key=key)

    def __len__(self) -> int:
        """
        Returns the length of the dataset (i.e. the number of molecules).

        :return: The length of the dataset.
        """
        return len(self.data)

    def __getitem__(self, item) -> Union[MoleculeDatapoint, List[MoleculeDatapoint]]:
        """
        Gets one or more MoleculeDatapoints via an index or slice.

        :param item: An index (int) or a slice object.
        :return: A MoleculeDatapoint if an int is provided or a list of MoleculeDatapoints if a slice is provided.
        """
        return self.data[item]


# Memoization
SMILES_TO_SENTENCE = {}

# add option 'radius'
def mol2sentence(smiles_batch: List[str], vocab, args: Namespace) -> List[dict]:
    output_list = []
    for smiles in smiles_batch:
        if smiles in SMILES_TO_SENTENCE:
            sentence = SMILES_TO_SENTENCE[smiles]
        else:
            mol = Chem.MolFromSmiles(smiles)
            if mol is not None:
                sentence = mol2alt_sentence(mol, radius=args.radius)
                SMILES_TO_SENTENCE[smiles] = sentence
            else:
                continue
        # convert to ids
        sentence = [vocab.stoi.get(token, vocab.unk_index) for i, token in enumerate(sentence)]
        sentence = [vocab.sos_index] + sentence + [vocab.eos_index]
        segment_label = ([1 for _ in range(len(sentence))])[:args.seq_len]

        input = sentence[:args.seq_len]
        padding = [vocab.pad_index for _ in range(args.seq_len - len(input))]
        input.extend(padding)
        segment_label.extend(padding)

        output = {
            'input': input,
            'segment_label': segment_label
        }

        output = {key: torch.tensor(value) for key, value in output.items()}
        output_list.append(output)

    return output_list



